#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <err.h>
#include <sysexits.h>
#include <elf.h>
#include <libelf.h>
#include <gelf.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/ptrace.h>
#include <string.h>
#include <errno.h>
#include <ctype.h>

#include "ftrace.h"
#include "list.h"
#include "my_nm.h"
#include "heap.h"
#include "getsymaddr.h"
#include "list_cb.h"

void			prepare_graph();
void			generate_dot_file();

static int	wait_singlestep(pid_t pid)
{
  int	status;

  if (-1 == wait4(pid, &status, 0, NULL))
  {
    perror("wait4");
    exit(1);
  }
  if (WIFEXITED(status) || WIFSIGNALED(status))
    return (0);
  else if (WIFSTOPPED(status))
    return (1);
  return (-1);
}

pid_t		trace_proc(pid_t pid, char **argv)
{
  if (pid == 0)
  {
    if (-1 == (pid = fork()))
    {
      perror("fork");
      return (-1);
    }
    if (pid)
      return (pid);
    if (-1 == ptrace(PT_TRACE_ME, pid, NULL, NULL))
    {
      perror("ptrace");
      return (-1);
    }
    execvp(argv[0], argv);
    perror("exec");
    exit(EXIT_FAILURE);
  }
  if (-1 == ptrace(PT_ATTACH, pid, NULL, NULL))
  {
    perror("ptrace");
    return (-1);
  }
  return (pid);
}

int	usage(char *bin)
{
  printf("Usage : %s [-p pid] | [bin]\n", bin);
  return (0);
}

int		start_ftrace(pid_t pid, char *filename)
{
  s_list	*symlist;
  s_heap	*heap;
  s_heap	*current;

  symlist = NULL;
  heap = NULL;
  current = NULL;
  if (filename && NULL == (symlist = my_nm(filename)))
    fprintf(stderr, "Unable to build symbol list\n");
  while (wait_singlestep(pid) > 0 &&
      (-1 != getsymaddr(pid, &heap, &current, symlist)))
    {
      if (-1 == ptrace(PT_STEP, pid, 0, 0))
	perror("ptrace");
    }
  prepare_graph();
  dsp_heap(heap, 0);
  generate_dot_file();
  destruct_heap(heap, &heap_dsc, NULL);
  destruct_symlist(symlist);
  return (0);
}

int	main(int argc, char *argv[])
{
  int		c;
  int		pflag;
  pid_t		pid;

  pid = 0;
  pflag = 0;
  while (EOF != (c = getopt(argc, argv, "+p:")))
    {
      if (c == 'p')
	{
	  pflag = 1;
	  pid = atoi(optarg);
	}
      else
	usage(argv[0]);
    }
  if ((optind >= argc && !pflag) || (optind < argc && pflag))
    return (usage(argv[0]));
  if (-1 == (pid = trace_proc(pid, &argv[optind])))
    return (1);
  return (start_ftrace(pid, argv[optind]));
}
